#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2021/9/14 12:10 上午
# @Author  : khan_long
# @Email   : longkehan15@qq.com
# @File    : dataset_reader.py
# @Software: PyCharm

import jsonlines
import re
import numpy as np
import spacy

regex_find_citation = re.compile(r"\(\s?(([A-Za-z\-]+\s)+([A-Za-z\-]+\.?)?,?\s\d{2,4}[a-c]?(;\s)?)+\s?\)|"
                                 r"\[(\d{1,3},\s?)+\d{1,3}\]|"
                                 r"\[[\d,-]+\]|(\([A-Z][a-z]+, \d+[a-c]?\))|"
                                 r"([A-Z][a-z]+ (et al\.)? \(\d+[a-c]?\))|"
                                 r"[A-Z][a-z]+ and [A-Z][a-z]+ \(\d+[a-c]?\)]")

# glove.6B.100d.txt.gz
spacy_tokenizer = spacy.load('en_core_web_sm')


def my_tokenizer(_sent):
    # TODO remove punctuation
    return [token.text for token in spacy_tokenizer(_sent)]


# class SectionTitleDataset(data.Dataset):
#     def __init__(self, text_field, label_field, _which, **kwargs):
#         fields = [("text", text_field), ("label", label_field)]
#
#         examples = []
#         if _which == 'acl':
#             for c_c, sec in acl_section_title_dataset_reader():
#                 examples.append(data.Example.fromlist([c_c, sec], fields))
#         else:
#             for c_c, sec in sci_section_title_dataset_reader():
#                 examples.append(data.Example.fromlist([c_c, sec], fields))
#         super(SectionTitleDataset, self).__init__(examples, fields, **kwargs)
#
#
# def acl_section_title_dataset_reader(_clean_citation=False):
#     for obj in jsonlines.open('./datasets/acl-arc/scaffolds/sections-scaffold-train.jsonl'):
#         citation_text = obj['text']
#
#         if _clean_citation:
#             citation_text = regex_find_citation.sub("", citation_text)
#
#         section_name = obj['section_name']
#
#         yield citation_text, section_name
#
#
def sci_section_title_dataset_reader(_clean_citation=False):
    for obj in jsonlines.open('./datasets/scicite/scaffolds/sections-scaffold-train.jsonl'):
        citation_text = obj['text']

        if _clean_citation:
            citation_text = regex_find_citation.sub("", citation_text)

        section_name = obj['section_name']

        yield citation_text, section_name


def acl_citation_worthiness_reader(_clean_citation=True):
    for obj in jsonlines.open('./datasets/acl-arc/scaffolds/cite-worthiness-scaffold-train.jsonl'):
        if _clean_citation:
            citation_text = obj['cleaned_cite_text']
        else:
            citation_text = obj['text']

        is_citation = obj['is_citation']

        yield citation_text, is_citation


#
# CITATION_TEXT_S = data.Field(sequential=True, tokenize=my_tokenizer, lower=True, fix_length=200)
# SECTION_LABEL = data.Field(sequential=False, use_vocab=False)
# acl_section_dataset = data.TabularDataset(
#     path='./datasets/acl-arc/scaffolds/sections-scaffold-train.jsonl',
#     format='json',
#     # key_in_json : (name_in_filed, filed obj)
#     fields={'text': ("text", CITATION_TEXT_S), 'section_name': ("section_label", SECTION_LABEL)}
# )
#
# CITATION_TEXT_W = data.Field(sequential=True, tokenize=my_tokenizer, lower=True, fix_length=200)
# CITATION_LABEL = data.Field(sequential=False, use_vocab=False)
# acl_worthiness_dataset = data.TabularDataset(
#     path='./datasets/acl-arc/scaffolds/cite-worthiness-scaffold-train.jsonl',
#     format='json',
#     fields={'text': ("text", CITATION_TEXT_W), 'is_citation': ("citation_label", CITATION_LABEL)}
# )

if __name__ == '__main__':
    from torchtext.vocab import GloVe

    glove = GloVe(name='6B', dim=100, cache='./.vector_cache')  # 与上面等价
    print(glove['token'])
